import json
import re
import numpy as np
import random
import torch

# import torch_npu

promptTemplates = json.load(open("./promptTemplates.json", "r", encoding="utf-8"))


# 设置随机数种子
def set_seed(seed=42):
    """固定 torch、numpy 和 random 的随机种子，确保实验可复现"""
    random.seed(seed)  # 固定 Python 内置的 random 模块
    np.random.seed(seed)  # 固定 numpy 的随机种子
    torch.manual_seed(seed)  # 固定 torch 的 CPU 随机种子
    torch.cuda.manual_seed(seed)  # 固定 torch 的 GPU 随机种子
    torch.cuda.manual_seed_all(seed)  # 如果使用多个 GPU，固定所有 GPU 的随机种子
    torch.backends.cudnn.deterministic = True  # 确保 cuDNN 计算是确定性的
    torch.backends.cudnn.benchmark = False  # 关闭 cuDNN 的自动优化，保证可复现性
    # torch_npu.npu.manual_seed(seed)
    # torch_npu.npu.manual_seed_all(seed)  # if using multiple NPUs


# ========== 评估准确率函数Accuracy==========
def compute_accuracy(prompts, completions, references, bench, id, issudoku):
    log = []
    cnt = 0
    for idx in range(len(id)):
        print(f"------[{bench}]    id: {id[idx]}---------------")
        correct = 0
        if not issudoku[idx]:
            try:
                clean_answer_ = completions[idx].split("{")[-1].split(":")[-1]
                if "," in clean_answer_:  # LogicVista有多选
                    clean_answer = ", ".join(re.findall("[ABCDEFGH]+", clean_answer_))
                else:
                    clean_answer = re.findall("[ABCDEFGH1234567890]+", clean_answer_)
                    if len(clean_answer) != 0:
                        clean_answer = clean_answer[0]
                        if clean_answer in references[idx].strip():
                            correct = 1
                            cnt += 1
                    else: # PuzzleVQA的有些选项是单词或词组
                        clean_answer = re.findall("[a-z ]+", clean_answer_)[0]
                        if clean_answer == references[idx].strip():
                            correct = 1
                            cnt += 1

            except Exception as e:
                pass
        else:  # 数独的判定
            try:
                clean_answer_ = re.findall("{.*}", completions[idx])[-1]
                clean_answer = "\n".join(list(json.loads(clean_answer_).values())[0].split())
                correct = 0
                if clean_answer == references[idx].strip():
                    correct = 1
                    cnt += 1
            except Exception as e:
                pass
        print(f"Prompts: {prompts[idx]}\nPrediction: {completions[idx]}\nReference: {references[idx]}\nCorrect: {correct}")
        log.append({"ID": id[idx], "Prediction": completions[idx], "Reference": references[idx], "Correct": correct})
    acc = cnt / len(references)
    print(f"Accuracy: {acc:.3f}")
    return {"Acc": acc, "Log": log}
